import os
from PyCmpltrtok.common import sep
import cv2 as cv


IMG_FORMATS = set(['bmp', 'dng', 'jpeg', 'jpg', 'mpo', 'png', 'tif', 'tiff', 'webp', 'pfm'])  # include image suffixes


class PicDirChecker(object):

    def __init__(self, dir):
        self.dir = dir
        self.names = os.listdir(dir)
        self.nf = len(self.names)  # number of files

    def __iter__(self):
        self.count = 0
        return self

    def __next__(self):

        def check_pic():
            if self.count >= self.nf:
                raise StopIteration
            name = self.names[self.count - 1]
            base, ext = os.path.splitext(name)
            ext = ext[1:].lower()
            if ext not in IMG_FORMATS:
                return None
            path = os.path.join(self.dir, name)
            if not os.path.isfile(path):
                return None
            img = cv.imread(path, cv.IMREAD_UNCHANGED)
            return name, path, img.shape

        while (True):
            self.count += 1
            res = check_pic()
            if res is not None:
                return res


if '__main__' == __name__:
    import pandas as pd

    pd.set_option('display.max_rows', None, 'display.max_columns', None, 'display.max_colwidth', 1000, 'display.expand_frame_repr', False)

    sep('check dir')
    checker = PicDirChecker('/home/asuspei/large_data/pic')
    res_list = []
    for idx, (name, path, shape) in enumerate(checker):
        print(idx, name, shape, path)
        res_list.append([idx, name.lower(), name, shape, path])
    sep('list to pandas')
    fr = pd.DataFrame(res_list)
    fr.sort_values([1], ascending=[True], inplace=True)
    fr2 = fr.copy()
    fr.reset_index(inplace=True)
    print(fr)
    sep('ordered by ori-index again')
    fr2.sort_index('index', ascending=[True], inplace=True)
    print(fr2)
    sep('All Over')
